
#include <apr_strings.h>

#include "../utils/Hash.h"
#include "../ssh/SSHClientWrapper.h"
#include "../ssh/SSHTelescope.h"
#include "../ssh/ShellChannel.h"
#include "../ssh/SCPSendChannel.h"
#include "WorkInstance.h"
#include "../iface/testerror.h"

#include "default_hosts.h"

/** Constructor for the work instance wrapper.
 * 
 * @param idbs Structure containing pointers to all the databases.
 * 
 * @param iname The name of the work instance.
 * 
 * @param iowner The owner of the work instance (parent). */
WorkInstance::WorkInstance(AllDBs *idbs, const char *iname, const char *iowner) : Actor(idbs->actorsdb)
{
   existing = 0;
   invalid = 0;
   dbs = idbs;
   if ((dbs) && (dbs->actorsdb->testActorExists(iname))) existing = 1;
   if (existing == 1)
      invalid = (dbs->actorsdb->isEquivalentPermission(iname, iowner) == 0);
   work_instance_id = strdup(iname);
   copied_owner = strdup(iowner);
   setName(iname);
   if (existing)
      getType();
   else
      setType(ACTOR_TYPE_WORKINSTANCE);
   setWorkInstanceDataStore(iowner);
}

/** Destructor for the work instance wrapper. Should commit
    changes. */
WorkInstance::~WorkInstance()
{
   if (work_instance_id) free(work_instance_id);
   work_instance_id = NULL;
   if (copied_owner) free(copied_owner);
   copied_owner = NULL;
}

/** Return a list of clusters which the work can in theory be run on,
 * given the clusters available to the top-level user.
 * 
 * @param clusters If this is not NULL, this is a comma-separated
 * limited list that the work instance is allowed to choose from.
 * 
 * @return The comma-separated list of cluster which the work instance
 * can run on. */
char *WorkInstance::getUpdatedClusterList(const char *clusters)
{
   std::vector<char *> cluster_list;
   if (clusters != NULL)
   {
      std::vector<char *> *unflat = StringUtils::unflattenStringList(clusters, ",");
      for (int i = 0; i < unflat->size(); i++)
      {
         cluster_list.push_back(strdup(unflat->at(i)));
      }
      StringUtils::destroyStringList(unflat);
   } else
   {
      char *ultimate_user = dbs->actorsdb->getUltimateUser(copied_owner);
      std::vector<char *> all_users_hosts;
      std::vector<int> all_users_ports;
      char *ultimate_ssh_user = StringUtils::varprintf("%s_ssh", ultimate_user);
      dbs->networkdb->getAllHostsAccessibleToUserID(ultimate_ssh_user, &all_users_hosts, &all_users_ports);
      char *hosts_copy = StringUtils::flattenStringList(&all_users_hosts, ",");
      errprint("user '%s' can access hosts '%s' for work", ultimate_user, hosts_copy);
      free(hosts_copy);
      free(ultimate_ssh_user);
      free(ultimate_user);
      for (int i = 0; i < all_users_hosts.size(); i++)
      {
         int host_is_node = 0;
         dbs->networkdb->sqlGetSingleIntf(&host_is_node, "select count(clustername) from 'nodes' where hostname='%s';", all_users_hosts.at(i));
         if (host_is_node)
         {
            char *cluster_name_for_host = NULL;
            dbs->networkdb->sqlGetSingleStringf(&cluster_name_for_host, "select clustername from 'nodes' where hostname='%s';", all_users_hosts.at(i));
            // check if it isn't already here
            int exists_in_list = 0;
            for (int j = 0; j < cluster_list.size(); j++)
               exists_in_list |= (strcmp(cluster_list.at(j), cluster_name_for_host) == 0);
            if (!(exists_in_list))
               cluster_list.push_back(cluster_name_for_host);
         }
      }
      StringUtils::destroyStringList(all_users_hosts);
   }
   char *final_string = StringUtils::flattenStringList(&cluster_list, ",");
   StringUtils::destroyStringList(cluster_list);
   return final_string;
}

/** Instantiate work based on the data store and instance information
 * given in the input message.
 * 
 * @param msgin The input message data, containing instance data.
 * 
 * @param dbs A stucture containing pointers to all the databases.
 * 
 * @param t_mp The memory pool from which to allocate.
 * 
 * @return Zero on success, non-zero on failure. */
int WorkInstance::instantiateWork(BrokerMessage &msgin, AllDBs *dbs, apr_pool_t *t_mp)
{
   const APRTable *headertable = msgin.getHeaderTable();
   // Get number of instances
   const char *instance_ids_string = headertable->get("Instances");
   int instance_count = 1;
   if ((instance_ids_string != NULL) && (strlen(instance_ids_string) > 0))
      instance_count = atoi(instance_ids_string);
   if (instance_count <= 0)
   {
      errprint("invalid instance count specified (%d)", instance_count);
      return -1;
   }
   for (int i = 0; i < instance_count; i++)
   {
      char *hex_id = Hash::makeHexadecimalID(32);
      WorkInstance t_instance(dbs, hex_id, msgin.getSender());
      free(hex_id);
      // probably want to start filling in the fields of the WorkInstance here before trying to start it!
      int retval = t_instance.tryStartingWork(msgin, t_mp, i);
      if (retval)
      {
         // if it's an invalid work instance, say so!
         t_instance.setUncommitable();
         return 1;
      } else
      {
         t_instance.setParent(msgin.getSender());
         char *instance_name = StringUtils::varprintf("instance%d", i+1);
         t_instance.setExternalAlias(instance_name);
         free(instance_name);
         t_instance.commit();
      }
   }
   return 0;
}

/** Try to start the work described in the input message.
 * 
 * @param msgin The input message containing instance and data store
 * information.
 * 
 * @param t_mp The memory pool from which to allocate.
 * 
 * @param i The instance number of this work instance.
 * 
 * @return Zero on success, non-zero on failure. */
int WorkInstance::tryStartingWork(BrokerMessage &msgin, apr_pool_t *t_mp, int i)
{
   const APRTable *headertable = msgin.getHeaderTable();
   const char *clusterids = headertable->get("Cluster-IDs");
   char *real_cluster_list_string = getUpdatedClusterList(clusterids);
   std::vector<char *> *real_cluster_list = StringUtils::unflattenStringList(real_cluster_list_string, ",");
   free(real_cluster_list_string);
   if ((real_cluster_list == NULL) || (real_cluster_list->size() == 0))
   {
      errprint("can't start work - no clusters available");
      if (real_cluster_list) delete real_cluster_list;
      return -1;
   }
   // Don't really have time to implement proper capabilities-based
   // selection & matching. Just select off the top of the list.
   char *cluster_choice = strdup(real_cluster_list->at(0));
   StringUtils::destroyStringList(real_cluster_list);
   // Find the hosts for the cluster
   errprint("cluster choice was '%s'", cluster_choice);
   std::vector<char *> *headnodenames = dbs->networkdb->sqlScanAllTextf("select hostname from 'nodes' where clustername='%s';", cluster_choice);
   if (headnodenames->size() == 0)
   {
      errprint("cannot get to cluster '%s', inaccessible or unknown cluster", cluster_choice);
      delete headnodenames;
      free(cluster_choice);
      return -1;
   }
   std::vector<char *> *elems = NULL;
   // maybe try to convert elems->at(2) to keyfingerprint binary?
   // should probably also have pendingdb
   SSHClient *sshcli = NULL;
   int aprsleep = 50000;
   int start_headnode = rand() % (headnodenames->size());
   // number of retries is either 2*headnodes or 5, whichever is larger (this way everything is double-checked).
   int retries = ((2*(headnodenames->size())) < 5) ? 5 : (2*(headnodenames->size()));
   while (retries--)
   {
      // Pick one at random
      const char *headnoderequested = (const char *)apr_pstrdup(t_mp, headnodenames->at((start_headnode + retries) % (headnodenames->size())));
      elems = dbs->networkdb->sqlScanAllTextf("select hostname,sshport,hostfingerprint from 'hosts' where hostname='%s';", headnoderequested);
      // SSH to the end - ShellChannel
      char *ult_user = newadb->getUltimateUser(copied_owner);
      char *ssh_username_special = StringUtils::mpprintf(t_mp, "%s_ssh", ult_user);
      SSHTelescope telescope(ssh_username_special, *(dbs->networkdb), *(dbs->keydb), elems->at(0), atoi(elems->at(1)), t_mp, headertable->get("Secret-Password"));
      free(ult_user);
      if (telescope.telescopeOkay())
      {
         sshcli = SSHClientWrapper::sshClientDefault(
            ((telescope.getLocalPort() == 0) ? (elems->at(0)) : "127.0.0.1"),
            ((telescope.getLocalPort() == 0) ? (atoi(elems->at(1))) : telescope.getLocalPort()), t_mp, NULL, 0,
            elems->at(0), atoi(elems->at(1)), NULL, ssh_username_special, dbs->networkdb, dbs->keydb, headertable->get("Secret-Password"));
         if (sshcli == NULL)
         {
            StringUtils::destroyStringList(elems);
            elems = NULL;
            errprint("SSH client was not able to start, will retry %d more times...", retries);
            apr_sleep(aprsleep);
            aprsleep <<= 1;
         }
      } else
      {
         errprint("telescope to host '%s:%s' failed", elems->at(0), elems->at(1));
         StringUtils::destroyStringList(elems);
         elems = NULL;
      }
      int schedulesuccess = 0;
      if (sshcli)
      {
         schedulesuccess = (scheduleWork(*sshcli, cluster_choice, t_mp, i, headertable->get("Secret-Password")) == 0);
         delete sshcli;
         if (schedulesuccess == 0)
            errprint("failed to schedule work via SSH to '%s'", headnoderequested);
      } else
      {
         errprint("SSH client was not able to start (was NULL)");
      }
      if (schedulesuccess) break;
   }
   StringUtils::destroyStringList(headnodenames);
   if (elems)
   {
      StringUtils::destroyStringList(elems);
      elems = NULL;
   }
   free(cluster_choice);
   return 0;
}

/** Get the directory location on the remote server.
 * 
 * @return A string containing the directory location, to be freed by
 * the caller. */
char *WorkInstance::getWorkDirectory(SSHClient &tclient, const char *workdir)
{
   char *tcommand = StringUtils::varprintf("sh -c 'echo -n \"%s\"'", workdir);
   ShellChannel tshell(tclient);
   APRFuture<char *> *rstringfuture = tshell.runCommandToStringFuture(tcommand);
   if (rstringfuture == NULL)
   {
      errprint("failed to get working directory");
      return NULL;
   }
   char *result_string = *(rstringfuture->get());
   delete rstringfuture;
   free(tcommand);
   return result_string;
}

/** Make the directory location recursively on the remote server.
 * 
 * @return A string containing the directory location, to be freed by
 * the caller. */
char *WorkInstance::makeWorkDirectory(SSHClient &tclient, const char *workdir)
{
   char *tcommand = StringUtils::varprintf("sh -c 'mkdir -p \"%s\"'", workdir);
   ShellChannel tshell(tclient);
   APRFuture<char *> *rstringfuture = tshell.runCommandToStringFuture(tcommand);
   if (rstringfuture == NULL)
   {
      errprint("failed to make working directory");
      return NULL;
   }
   char *result_string = *(rstringfuture->get());
   delete rstringfuture;
   free(tcommand);
   return result_string;
}

/** Securely copy the given key to the remote file given.
 * 
 * @param sshcli Open SSH client connection to the remote machine.
 * 
 * @param remotekeyfilename The filename to write the key to.
 * 
 * @param keyname The name of the key locally to use.
 * 
 * @param keytype The type of the key locally to use.
 * 
 * @param t_mp The memory pool to allocate from.
 * 
 * @return Zero on success, non-zero on failure. */
int WorkInstance::remoteKeyCopy(SSHClient &sshcli, const char *remotekeyfilename, const char *keyname, Key::keypurposetype keytype, apr_pool_t *t_mp)
{
   char *tk_kdata = NULL;
   Key t_key(dbs->keydb, keyname, keytype);
   const KeyTableView *tableview = NULL;
   switch (keytype)
   {
      case Key::KEYPURPOSE_ENCRYPTION: tableview = dbs->keydb->encryptionTable(); break;
      case Key::KEYPURPOSE_DECRYPTION: tableview = dbs->keydb->decryptionTable(); break;
      case Key::KEYPURPOSE_SIGNING: tableview = dbs->keydb->signingTable(); break;
      case Key::KEYPURPOSE_CHECKING: tableview = dbs->keydb->checkingTable(); break;
      default: break;
   }
   if (tableview == NULL)
   {
      errprint("unknown key type for '%s'", keyname);
      return -1;
   }
   AbstractKey *abskey = t_key.getKey(tableview->getWIPPassword());
   abskey->exportKeyToASCII(&tk_kdata);
   delete abskey;
   if (tk_kdata)
   {
      SCPSendChannel schan(sshcli);
      APRFuture<int> *future_result = schan.copyBufferToRemoteFileFuture(remotekeyfilename, tk_kdata, strlen(tk_kdata), 0700);
      if (future_result == NULL)
      {
         errprint("failed to get buffer to copy to remote file '%s'", remotekeyfilename);
         return -1;
      }
      int result = *(future_result->get());
      delete future_result;
      if (result != 0)
      {
         errprint("copy failed for plain key going to '%s'", remotekeyfilename);
         return -1;
      }
      free(tk_kdata);
   }
   return 0;
}

/** Copy the keys and run the scripts remotely needed to schedule the
 * work.
 * 
 * @param sshcli Open SSH client connection to the remote machine.
 * 
 * @param clustername The name of the cluster used.
 * 
 * @param t_mp The memory pool to allocate from.
 * 
 * @param i The local work instance number.
 * 
 * @param secretpass The secret SSH key chain unlocking password to
 * resubmit further workpackets.
 * 
 * @return Zero on success, one on failure. */
int WorkInstance::scheduleWork(SSHClient &sshcli, const char *clustername, apr_pool_t *t_mp, int i, const char *secretpass)
{
   if ((clustername == NULL) || (strlen(clustername) == 0))
   {
      errprint("Can't have NULL/empty cluster name");
      return 1;
   }
   // we probably should issue the creation of a instance_<uuid>
   // directory, so that if this is writing into the same space
   // clobbering doesn't happen.
   char *inp_key_name = StringUtils::mpprintf(t_mp, "%s_inp", getWorkInstanceDataStore());
   char *out_key_name = StringUtils::mpprintf(t_mp, "%s_out", getWorkInstanceDataStore());
   char *working_dir = NULL;
   dbs->networkdb->sqlGetSingleStringf(&working_dir, "select writedir from 'clusters' where clustername='%s';", clustername);
   char *workdir = NULL;
   if (working_dir)
   {
      workdir = getWorkDirectory(sshcli, working_dir);
      free(working_dir);
   }
   if (workdir == NULL)
   {
      errprint("unset working directory for cluster '%s', setting to '/tmp'", clustername);
      workdir = strdup("/tmp");
   }
   char *workdir2 = StringUtils::mpprintf(t_mp, "%s/%s", workdir, getName());
   free(workdir);
   workdir = NULL;
   char *resstr = makeWorkDirectory(sshcli, workdir2);
   if ((resstr != NULL) && (strlen(resstr) > 0))
      errprint("mkdir returned '%s'", resstr);
   if (resstr) free(resstr);
   char *tk_keyname = NULL;
   tk_keyname = StringUtils::mpprintf(t_mp, "%s/%s.enc", workdir2, getWorkInstanceDataStore());
   remoteKeyCopy(sshcli, tk_keyname, getWorkInstanceDataStore(), Key::KEYPURPOSE_ENCRYPTION, t_mp);
   tk_keyname = StringUtils::mpprintf(t_mp, "%s/%s.dec", workdir2, getWorkInstanceDataStore());
   remoteKeyCopy(sshcli, tk_keyname, getWorkInstanceDataStore(), Key::KEYPURPOSE_DECRYPTION, t_mp);
   tk_keyname = StringUtils::mpprintf(t_mp, "%s/%s.sgn", workdir2, getWorkInstanceDataStore());
   remoteKeyCopy(sshcli, tk_keyname, getWorkInstanceDataStore(), Key::KEYPURPOSE_SIGNING, t_mp);
   tk_keyname = StringUtils::mpprintf(t_mp, "%s/%s.chk", workdir2, getWorkInstanceDataStore());
   remoteKeyCopy(sshcli, tk_keyname, getWorkInstanceDataStore(), Key::KEYPURPOSE_CHECKING, t_mp);
   tk_keyname = StringUtils::mpprintf(t_mp, "%s/%s.dec", workdir2, inp_key_name);
   remoteKeyCopy(sshcli, tk_keyname, inp_key_name, Key::KEYPURPOSE_DECRYPTION, t_mp);
   tk_keyname = StringUtils::mpprintf(t_mp, "%s/%s.enc", workdir2, out_key_name);
   remoteKeyCopy(sshcli, tk_keyname, out_key_name, Key::KEYPURPOSE_ENCRYPTION, t_mp);
   tk_keyname = StringUtils::mpprintf(t_mp, "%s/%s.enc", workdir2, inp_key_name);
   remoteKeyCopy(sshcli, tk_keyname, inp_key_name, Key::KEYPURPOSE_ENCRYPTION, t_mp);
   // All six, whew!
   if (secretpass)
   {
      // Copy the encrypted password
      tk_keyname = StringUtils::mpprintf(t_mp, "%s/password", workdir2);
      EncryptWrapper *encwrap = dbs->keydb->getEncrypter(getWorkInstanceDataStore(), NULL, 0);
      apr_size_t newlength = 0;
      char *encd_pass = encwrap->encryptDataBuffer(secretpass, (apr_size_t)strlen(secretpass), &newlength);
      delete encwrap;
      SCPSendChannel schan(sshcli);
      APRFuture<int> *future_result = schan.copyBufferToRemoteFileFuture(tk_keyname, encd_pass, newlength, 0700);
      if (future_result == NULL)
      {
         errprint("fatal - could not copy key buffer to remote file...");
         return 1;
      }
      int result = *(future_result->get());
      delete future_result;
      free(encd_pass);
      if (result != 0)
      {
         errprint("copy failed for encrypted password going to '%s'", tk_keyname);
         return 1;
      }
   }
   // All six, whew!
   // now we chain the new versions of our cluster scripts
   const char *remote_script = (const char *)StringUtils::mpprintf(t_mp, "export CONSPIRE_DEMAND_JOB=%s; "
                                                                         "export CONSPIRE_WRITEABLEDIR=%s; "
                                                                         "export CONSPIRE_KEYID=%s; "
                                                                         "export CONSPIRE_CLUSTERID=%s; "
                                                                         "export CONSPIRE_TMAX=%" APR_INT64_T_FMT "; "
                                                                         "(cd \"$CONSPIRE_WRITEABLEDIR\"; "
                                                                            "if [ `which wget` ]; then "
                                                                               "wget http://%s:%d/hosted_files/queue.sh -O queue_$CONSPIRE_KEYID.sh; "
                                                                            "else "
                                                                               "curl http://%s:%d/hosted_files/queue.sh > queue_$CONSPIRE_KEYID.sh; "
                                                                            "fi; "
                                                                         "chmod +x queue_$CONSPIRE_KEYID.sh; "
                                                                         "sh queue_$CONSPIRE_KEYID.sh; "
                                                                         "rm queue_$CONSPIRE_KEYID.sh)",
      getName(), workdir2, getWorkInstanceDataStore(), clustername, getRemainingLife(),
      DEFAULT_HOST, DEFAULT_PORT, DEFAULT_HOST, DEFAULT_PORT);
   ShellChannel tshell(sshcli);
   APRFuture<char *> *rstringfuture = tshell.runCommandToStringFuture(remote_script);
   if (rstringfuture == NULL)
   {
      errprint("fatal - failed to remotely execute queueing script");
      return 1;
   }
   char *result_string = *(rstringfuture->get());
   delete rstringfuture;
   if ((result_string != NULL) && (strlen(result_string) > 0))
   {
      // probably want to filter this to remove whitespace
      const char *trimmed_string = StringUtils::mpTrimWhitespaceFromBoth(result_string, t_mp);
      errprint("remote end for cluster '%s' returned message (which should be the job ID):\n%s", clustername, trimmed_string);
      setWorkInstanceCompute(trimmed_string);
      setWorkInstanceAccessCluster(clustername);
      setWorkInstanceLocalID(i);
   }
   if (result_string) free(result_string);
   return 0;
}
